#include "decodec.h"

#define cxx_log(fmt,...)\
            do {\
                printf(fmt,##__VA_ARGS__);\
            } while(0)

decodec::decodec() : mFps(0),
                     mEos(0),
                     mDisPlay(0),
                     mFrmCnt(0),
                     mFrmRate(0.0),
                     mFin(NULL),
                     mFout(NULL),
                     mPktBuf(NULL),
                     mCtx(NULL),
                     mApi(NULL),
                     mPkt(NULL),
                     mBuffer(NULL),
                     mFrmGrp(NULL)
{

}

decodec::~decodec() {
    if(mFin) {
        fclose(mFin);
        mFin = NULL;
    }

    if(mFout) {
        fclose(mFout);
        mFout = NULL;
    }

    if(mPktBuf) {
        delete[] mPktBuf;
        mPktBuf = NULL;
    }

    if (mCtx) {
        mpp_destroy(mCtx);
        mCtx = NULL;
    }

    if (mPkt) {
        mpp_packet_deinit(&mPkt);
        mPkt = NULL;
    }

    if (mBuffer) {
        mpp_buffer_put(mBuffer);
        mBuffer = NULL;
    }

    if (mFrmGrp) {
        mpp_buffer_group_put(mFrmGrp);
        mFrmGrp = NULL;
    }
}

int decodec::init(void *args) {
    int ret = 0;
    FileSource *arg = (FileSource *)args;
    RK_U32 need_split = 1;
    MppParam param = NULL;
    const char *file_input = arg->input.c_str();
    const char *file_output = arg->output.c_str();

    mDisPlay = arg->display;
    srcW = arg->width;
    srcH = arg->height;
    dstW = arg->width;
    dstH = arg->height;
    
    mFin = fopen(file_input, "rb");
    if (!mFin) {
        cxx_log("failed to open input file %s.\n", file_input);
        return -1;
    }

    mFout = fopen(file_output, "wb+");
    if (!mFout) {
        cxx_log("failed to open output file %s.\n", file_output);
        return -2;
    }

    mPktBuf = new char[SZ_4K];
    if (!mPktBuf) {
        cxx_log("failed to malloc mPktBuf.\n");
        return -3;
    }

    ret = mpp_packet_init(&mPkt, mPktBuf, SZ_4K);
    if (ret) {
        cxx_log("failed to exec mpp_packet_init ret %d", ret);
        return -4;
    }

    ret = mpp_create(&mCtx, &mApi);
    if (ret != MPP_OK) {
        cxx_log("failed to exec mpp_create ret %d", ret);
        return -5;
    }

    param = &need_split;
    ret = mApi->control(mCtx, MPP_DEC_SET_PARSER_SPLIT_MODE, param);
    if (ret) {
        cxx_log("failed to control MPP_DEC_SET_PARSER_SPLIT_MODE ret %d\n", ret);
        return  -6;
    }

    ret = mpp_init(mCtx, MPP_CTX_DEC, arg->type);
    if (ret != MPP_OK) {
        cxx_log("failed to exec mpp_init ret %d\n", ret);
        return -7;
    }

    return 0;
}

RK_S64 get_time()
{
    struct timespec time = {0, 0};
    clock_gettime(CLOCK_MONOTONIC, &time);
    return (RK_S64)time.tv_sec * 1000000 + (RK_S64)time.tv_nsec / 1000;
}

int decodec::decode(){
    int ret = 0;

    mTimeS = get_time();

    while(!mEos) {
        size_t file_input_size = fread(mPktBuf, 1, SZ_4K, mFin);
        if(file_input_size != SZ_4K || feof(mFin)) {
            cxx_log("%p found last packet\n", mCtx);
            mEos = 1;
        }

        mpp_packet_write(mPkt, 0, mPktBuf, file_input_size);
        mpp_packet_set_pos(mPkt, mPktBuf);
        mpp_packet_set_length(mPkt, file_input_size);
        if(mEos) {
            mpp_packet_set_eos(mPkt);
        }
        
        do {
            RK_U32 pkt_done = 0;
            RK_U32 frm_eos = 0;
            RK_S32 times = 30;

            if (!pkt_done) {
                ret = mApi->decode_put_packet(mCtx, mPkt);
                if (MPP_OK == ret) {
                    pkt_done = 1;
                }
            }

            do {
                RK_S32 get_frm = 0;
                MppFrame frame = NULL;

            try_again:
                ret = mApi->decode_get_frame(mCtx, &frame);
                if (MPP_ERR_TIMEOUT == ret) {
                    if (times > 0) {
                        times--;
                        msleep(1);
                        goto try_again;
                    }
                    cxx_log("%p decode_get_frame failed too much time\n", mCtx);
                }
                if (ret) {
                    cxx_log("%p decode_get_frame failed ret %d\n", mCtx, ret);
                    break;
                }

                if (frame) {
                    if (mpp_frame_get_info_change(frame)) {
                        RK_U32 width = mpp_frame_get_width(frame);
                        RK_U32 height = mpp_frame_get_height(frame);
                        RK_U32 hor_stride = mpp_frame_get_hor_stride(frame);
                        RK_U32 ver_stride = mpp_frame_get_ver_stride(frame);
                        RK_U32 buf_size = mpp_frame_get_buf_size(frame);

                        cxx_log("%p decode_get_frame get info changed found\n", mCtx);
                        cxx_log("%p decoder require buffer w:h [%d:%d] stride [%d:%d] buf_size %d\n",
                                mCtx, width, height, hor_stride, ver_stride, buf_size);

                        ret = mpp_buffer_group_get_internal(&mFrmGrp, MPP_BUFFER_TYPE_DMA_HEAP);
                        if (ret) {
                            cxx_log("get mpp buffer group failed ret %d.\n", ret);
                            break;
                        }
                        
                        ret = mApi->control(mCtx, MPP_DEC_SET_EXT_BUF_GROUP, mFrmGrp);
                        if (ret) {
                            cxx_log("%p set buffer group failed ret %d\n", mCtx, ret);
                            break;
                        }
                        
                        ret = mpp_buffer_group_limit_config(mFrmGrp, buf_size, 24);
                        if (ret) {
                            cxx_log("limit buffer group failed ret %d\n", ret);
                            break;
                        }

                        ret = mApi->control(mCtx, MPP_DEC_SET_INFO_CHANGE_READY, NULL);
                        if (ret) {
                            cxx_log("%p info change ready failed ret %d\n", mCtx, ret);
                            break;
                        }
                    } else {
                        char log_buf[256];
                        RK_S32 log_size = sizeof(log_buf) - 1;
                        RK_S32 log_len = 0;
                        RK_U32 err_info = mpp_frame_get_errinfo(frame);
                        RK_U32 discard = mpp_frame_get_discard(frame);

                        log_len += snprintf(log_buf + log_len, log_size - log_len,
                                            "decode get frame %d", mFrmCnt);

                        if (mpp_frame_has_meta(frame)) {
                            MppMeta meta = mpp_frame_get_meta(frame);
                            RK_S32 temporal_id = 0;

                            mpp_meta_get_s32(meta, KEY_TEMPORAL_ID, &temporal_id);

                            log_len += snprintf(log_buf + log_len, log_size - log_len,
                                                " tid %d", temporal_id);
                        }

                        if (err_info || discard) {
                            log_len += snprintf(log_buf + log_len, log_size - log_len,
                                                " err %x discard %x", err_info, discard);
                        }
                        cxx_log("%p %s\n", mCtx, log_buf);

                        mFrmCnt++;
                        if (mFout && !err_info)
                            dump_mpp_frame_to_file(frame, mFout);
                    }
                    frm_eos = mpp_frame_get_eos(frame);
                    mpp_frame_deinit(&frame);
                    get_frm = 1;
                }

                // try get runtime frame memory usage
                if (mFrmGrp) {
                    size_t usage = mpp_buffer_group_usage(mFrmGrp);
                    if (usage > max_usage)
                        max_usage = usage;
                }

                // if last packet is send but last frame is not found continue
                if (mEos && pkt_done && !frm_eos) {
                    msleep(1);
                    continue;
                }

                if (frm_eos) {
                    cxx_log("%p found last frame\n", mCtx);
                    break;
                }

                if (get_frm)
                    continue;
                break;
            } while (1);

            if (pkt_done)
                break;

            msleep(1);
        } while (1);
    }

    mTimeE = get_time();
    mTimeDiff = (mTimeE - mTimeS) / 1000;
    mFrmRate = (float)(mFrmCnt * 1000) / mTimeDiff;

    cxx_log("decode %d frames time %lld ms fps %3.2f\n",
        mFrmCnt, mTimeDiff, mFrmRate);
    cxx_log("test success max memory %.2f MB\n", max_usage / (float)(1 << 20));
    return ret;
}

static void rearrange_pix(RK_U8 *tmp_line, RK_U8 *base, RK_U32 n)
{
    RK_U16 * pix = (RK_U16 *)(tmp_line + n * 16);
    RK_U16 * base_u16 = (RK_U16 *)(base + n * 10);

    pix[0] =  base_u16[0] & 0x03FF;
    pix[1] = (base_u16[0] & 0xFC00) >> 10 | (base_u16[1] & 0x000F) << 6;
    pix[2] = (base_u16[1] & 0x3FF0) >> 4;
    pix[3] = (base_u16[1] & 0xC000) >> 14 | (base_u16[2] & 0x00FF) << 2;
    pix[4] = (base_u16[2] & 0xFF00) >> 8  | (base_u16[3] & 0x0003) << 8;
    pix[5] = (base_u16[3] & 0x0FFC) >> 2;
    pix[6] = (base_u16[3] & 0xF000) >> 12 | (base_u16[4] & 0x003F) << 4;
    pix[7] = (base_u16[4] & 0xFFC0) >> 6;
}

void decodec::dump_mpp_frame_to_file(MppFrame frame, FILE *fp)
{
    RK_U32 width    = 0;
    RK_U32 height   = 0;
    RK_U32 h_stride = 0;
    RK_U32 v_stride = 0;
    MppFrameFormat fmt  = MPP_FMT_YUV420SP;
    MppBuffer buffer    = NULL;
    RK_U8 *base = NULL;

    if (NULL == fp || NULL == frame)
        return ;

    width    = mpp_frame_get_width(frame);
    height   = mpp_frame_get_height(frame);
    h_stride = mpp_frame_get_hor_stride(frame);
    v_stride = mpp_frame_get_ver_stride(frame);
    fmt      = mpp_frame_get_fmt(frame);
    buffer   = mpp_frame_get_buffer(frame);

    if (NULL == buffer)
        return ;

    base = (RK_U8 *)mpp_buffer_get_ptr(buffer);

    if (MPP_FRAME_FMT_IS_FBC(fmt) || MPP_FRAME_FMT_IS_RKFBC(fmt)) {
        RK_S32 buf_size = mpp_buffer_get_size(buffer);

        fwrite(base, 1, buf_size, fp);
        return ;
    }

    if (MPP_FRAME_FMT_IS_RGB(fmt) && MPP_FRAME_FMT_IS_LE(fmt)) {
        fmt = (MppFrameFormat)(fmt & MPP_FRAME_FMT_MASK);;
    }
    switch (fmt & MPP_FRAME_FMT_MASK) {
    case MPP_FMT_YUV422SP : {
        /* YUV422SP -> YUV422P for better display */
        RK_U32 i, j;
        RK_U8 *base_y = base;
        RK_U8 *base_c = base + h_stride * v_stride;
        RK_U8 *tmp = mpp_malloc(RK_U8, h_stride * height * 2);
        RK_U8 *tmp_u = tmp;
        RK_U8 *tmp_v = tmp + width * height / 2;

        for (i = 0; i < height; i++, base_y += h_stride)
            fwrite(base_y, 1, width, fp);

        for (i = 0; i < height; i++, base_c += h_stride) {
            for (j = 0; j < width / 2; j++) {
                tmp_u[j] = base_c[2 * j + 0];
                tmp_v[j] = base_c[2 * j + 1];
            }
            tmp_u += width / 2;
            tmp_v += width / 2;
        }

        fwrite(tmp, 1, width * height, fp);
        mpp_free(tmp);
    } break;
    case MPP_FMT_YUV420SP_VU :
    case MPP_FMT_YUV420SP : {
        RK_U32 i;
        RK_U8 *base_y = base;
        RK_U8 *base_c = base + h_stride * v_stride;

        for (i = 0; i < height; i++, base_y += h_stride) {
            fwrite(base_y, 1, width, fp);
        }
        for (i = 0; i < height / 2; i++, base_c += h_stride) {
            fwrite(base_c, 1, width, fp);
        }
    } break;
    case MPP_FMT_YUV420P : {
        RK_U32 i;
        RK_U8 *base_y = base;
        RK_U8 *base_c = base + h_stride * v_stride;

        for (i = 0; i < height; i++, base_y += h_stride) {
            fwrite(base_y, 1, width, fp);
        }
        for (i = 0; i < height / 2; i++, base_c += h_stride / 2) {
            fwrite(base_c, 1, width / 2, fp);
        }
        for (i = 0; i < height / 2; i++, base_c += h_stride / 2) {
            fwrite(base_c, 1, width / 2, fp);
        }
    } break;
    case MPP_FMT_YUV420SP_10BIT : {
        RK_U32 i, k;
        RK_U8 *base_y = base;
        RK_U8 *base_c = base + h_stride * v_stride;
        RK_U8 *tmp_line = (RK_U8 *)mpp_malloc(RK_U16, width);

        if (!tmp_line) {
            cxx_log("tmp_line malloc fail");
            return;
        }

        for (i = 0; i < height; i++, base_y += h_stride) {
            for (k = 0; k < MPP_ALIGN(width, 8) / 8; k++)
                rearrange_pix(tmp_line, base_y, k);
            fwrite(tmp_line, width * sizeof(RK_U16), 1, fp);
        }

        for (i = 0; i < height / 2; i++, base_c += h_stride) {
            for (k = 0; k < MPP_ALIGN(width, 8) / 8; k++)
                rearrange_pix(tmp_line, base_c, k);
            fwrite(tmp_line, width * sizeof(RK_U16), 1, fp);
        }

        MPP_FREE(tmp_line);
    } break;
    case MPP_FMT_YUV444SP : {
        /* YUV444SP -> YUV444P for better display */
        RK_U32 i, j;
        RK_U8 *base_y = base;
        RK_U8 *base_c = base + h_stride * v_stride;
        RK_U8 *tmp = mpp_malloc(RK_U8, h_stride * height * 2);
        RK_U8 *tmp_u = tmp;
        RK_U8 *tmp_v = tmp + width * height;

        for (i = 0; i < height; i++, base_y += h_stride)
            fwrite(base_y, 1, width, fp);

        for (i = 0; i < height; i++, base_c += h_stride * 2) {
            for (j = 0; j < width; j++) {
                tmp_u[j] = base_c[2 * j + 0];
                tmp_v[j] = base_c[2 * j + 1];
            }
            tmp_u += width;
            tmp_v += width;
        }

        fwrite(tmp, 1, width * height * 2, fp);
        mpp_free(tmp);
    } break;
    case MPP_FMT_YUV444SP_10BIT : {
        RK_U32 i, k;
        RK_U8 *base_y = base;
        RK_U8 *base_c = base + h_stride * v_stride;
        RK_U8 *tmp_line = (RK_U8 *)mpp_malloc(RK_U16, width);

        if (!tmp_line) {
            cxx_log("tmp_line malloc fail");
            return;
        }

        for (i = 0; i < height; i++, base_y += h_stride) {
            for (k = 0; k < MPP_ALIGN(width, 8) / 8; k++)
                rearrange_pix(tmp_line, base_y, k);
            fwrite(tmp_line, width * sizeof(RK_U16), 1, fp);
        }

        for (i = 0; i < (height * 2); i++, base_c += h_stride) {
            for (k = 0; k < MPP_ALIGN(width, 8) / 8; k++)
                rearrange_pix(tmp_line, base_c, k);
            fwrite(tmp_line, width * sizeof(RK_U16), 1, fp);
        }

        MPP_FREE(tmp_line);
    } break;
    case MPP_FMT_YUV400: {
        RK_U32 i;
        RK_U8 *base_y = base;
        RK_U8 *tmp = mpp_malloc(RK_U8, h_stride * height);

        for (i = 0; i < height; i++, base_y += h_stride)
            fwrite(base_y, 1, width, fp);

        mpp_free(tmp);
    } break;
    case MPP_FMT_ARGB8888:
    case MPP_FMT_ABGR8888:
    case MPP_FMT_BGRA8888:
    case MPP_FMT_RGBA8888: {
        RK_U32 i;
        RK_U8 *base_y = base;
        RK_U8 *tmp = mpp_malloc(RK_U8, width * height * 4);

        for (i = 0; i < height; i++, base_y += h_stride)
            fwrite(base_y, 1, width * 4, fp);

        mpp_free(tmp);
    } break;
    case MPP_FMT_YUV422_YUYV:
    case MPP_FMT_YUV422_YVYU:
    case MPP_FMT_YUV422_UYVY:
    case MPP_FMT_YUV422_VYUY:
    case MPP_FMT_RGB565:
    case MPP_FMT_BGR565:
    case MPP_FMT_RGB555:
    case MPP_FMT_BGR555:
    case MPP_FMT_RGB444:
    case MPP_FMT_BGR444: {
        RK_U32 i;
        RK_U8 *base_y = base;
        RK_U8 *tmp = mpp_malloc(RK_U8, width * height * 2);

        for (i = 0; i < height; i++, base_y += h_stride)
            fwrite(base_y, 1, width * 2, fp);

        mpp_free(tmp);
    } break;
    case MPP_FMT_RGB888: {
        RK_U32 i;
        RK_U8 *base_y = base;
        RK_U8 *tmp = mpp_malloc(RK_U8, width * height * 3);

        for (i = 0; i < height; i++, base_y += h_stride)
            fwrite(base_y, 1, width * 3, fp);

        mpp_free(tmp);
    } break;
    default : {
        cxx_log("not supported format %d\n", fmt);
    } break;
    }
}

int decodec::deinit() {
    int ret = 0;

    ret = mApi->reset(mCtx);
    if (ret != MPP_OK) {
        cxx_log("failed to exec mApi->reset\n");
        return -1;
    }

    return 0;
}